import { useChainCallback } from '@/composables/functional/useChainCallback'
import { NodeSlotType } from '@/lib/litegraph/src/types/globalEnums'
import type {
  ISlotType,
  INodeInputSlot,
  INodeOutputSlot
} from '@/lib/litegraph/src/interfaces'
import type { LGraphNode } from '@/lib/litegraph/src/LGraphNode'
import { LiteGraph } from '@/lib/litegraph/src/litegraph'
import type { LLink } from '@/lib/litegraph/src/LLink'
import { commonType } from '@/lib/litegraph/src/utils/type'
import { transformInputSpecV1ToV2 } from '@/schemas/nodeDef/migration'
import type { ComboInputSpec, InputSpec } from '@/schemas/nodeDefSchema'
import type { InputSpec as InputSpecV2 } from '@/schemas/nodeDef/nodeDefSchemaV2'
import {
  zAutogrowOptions,
  zDynamicComboInputSpec
} from '@/schemas/nodeDefSchema'
import { useLitegraphService } from '@/services/litegraphService'
import { app } from '@/scripts/app'
import type { ComfyApp } from '@/scripts/app'

const INLINE_INPUTS = false

type MatchTypeNode = LGraphNode &
  Pick<Required<LGraphNode>, 'comfyMatchType' | 'onConnectionsChange'>

function ensureWidgetForInput(node: LGraphNode, input: INodeInputSlot) {
  if (input.widget?.name) return
  node.widgets ??= []
  node.widgets.push({
    name: input.name,
    y: 0,
    type: 'shim',
    options: {},
    draw(ctx, _n, _w, y) {
      ctx.save()
      ctx.fillStyle = LiteGraph.NODE_TEXT_COLOR
      ctx.fillText(input.label ?? input.name, 20, y + 15)
      ctx.restore()
    }
  })
  input.alwaysVisible = true
  input.widget = { name: input.name }
}

function dynamicComboWidget(
  node: LGraphNode,
  inputName: string,
  untypedInputData: InputSpec,
  appArg: ComfyApp,
  widgetName?: string
) {
  const { addNodeInput } = useLitegraphService()
  const parseResult = zDynamicComboInputSpec.safeParse(untypedInputData)
  if (!parseResult.success) throw new Error('invalid DynamicCombo spec')
  const inputData = parseResult.data
  const options = Object.fromEntries(
    inputData[1].options.map(({ key, inputs }) => [key, inputs])
  )
  const subSpec: ComboInputSpec = [Object.keys(options), {}]
  const { widget, minWidth, minHeight } = app.widgets['COMBO'](
    node,
    inputName,
    subSpec,
    appArg,
    widgetName
  )
  let currentDynamicNames: string[] = []
  const updateWidgets = (value?: string) => {
    if (!node.widgets) throw new Error('Not Reachable')
    const newSpec = value ? options[value] : undefined
    const inputsToRemove: Record<string, INodeInputSlot> = {}
    for (const name of currentDynamicNames) {
      const input = node.inputs.find((input) => input.name === name)
      if (input) inputsToRemove[input.name] = input
      const widgetIndex = node.widgets.findIndex(
        (widget) => widget.name === name
      )
      if (widgetIndex === -1) continue
      node.widgets[widgetIndex].value = undefined
      node.widgets.splice(widgetIndex, 1)
    }
    currentDynamicNames = []
    if (!newSpec) {
      for (const input of Object.values(inputsToRemove)) {
        const inputIndex = node.inputs.findIndex((inp) => inp === input)
        if (inputIndex === -1) continue
        node.removeInput(inputIndex)
      }
      return
    }

    const insertionPoint = node.widgets.findIndex((w) => w === widget) + 1
    const startingLength = node.widgets.length
    const initialInputIndex =
      node.inputs.findIndex((i) => i.name === widget.name) + 1
    let startingInputLength = node.inputs.length
    if (insertionPoint === 0)
      throw new Error("Dynamic widget doesn't exist on node")
    const inputTypes: [Record<string, InputSpec> | undefined, boolean][] = [
      [newSpec.required, false],
      [newSpec.optional, true]
    ]
    for (const [inputType, isOptional] of inputTypes)
      for (const key in inputType ?? {}) {
        const name = `${widget.name}.${key}`
        const specToAdd = transformInputSpecV1ToV2(inputType![key], {
          name,
          isOptional
        })
        specToAdd.display_name = key
        addNodeInput(node, specToAdd)
        currentDynamicNames.push(name)
        if (INLINE_INPUTS) ensureWidgetForInput(node, node.inputs.at(-1)!)
        if (
          !inputsToRemove[name] ||
          Array.isArray(inputType![key][0]) ||
          !LiteGraph.isValidConnection(
            inputsToRemove[name].type,
            inputType![key][0]
          )
        )
          continue
        node.inputs.at(-1)!.link = inputsToRemove[name].link
        inputsToRemove[name].link = null
      }

    for (const input of Object.values(inputsToRemove)) {
      const inputIndex = node.inputs.findIndex((inp) => inp === input)
      if (inputIndex === -1) continue
      if (inputIndex < initialInputIndex) startingInputLength--
      node.removeInput(inputIndex)
    }
    const inputInsertionPoint =
      node.inputs.findIndex((i) => i.name === widget.name) + 1
    const addedWidgets = node.widgets.splice(startingLength)
    node.widgets.splice(insertionPoint, 0, ...addedWidgets)
    if (inputInsertionPoint === 0) {
      if (
        addedWidgets.length === 0 &&
        node.inputs.length !== startingInputLength
      )
        //input is inputOnly, but lacks an insertion point
        throw new Error('Failed to find input socket for ' + widget.name)
      return
    }
    const addedInputs = spliceInputs(node, startingInputLength).map(
      (addedInput) => {
        const existingInput = node.inputs.findIndex(
          (existingInput) => addedInput.name === existingInput.name
        )
        return existingInput === -1
          ? addedInput
          : spliceInputs(node, existingInput, 1)[0]
      }
    )
    //assume existing inputs are in correct order
    spliceInputs(node, inputInsertionPoint, 0, ...addedInputs)
    node.size[1] = node.computeSize([...node.size])[1]
    if (!node.graph) return
    node._setConcreteSlots()
    node.arrange()
    app.canvas?.setDirty(true, true)
  }
  //A little hacky, but onConfigure won't work.
  //It fires too late and is overly disruptive
  let widgetValue = widget.value
  Object.defineProperty(widget, 'value', {
    get() {
      return widgetValue
    },
    set(value) {
      widgetValue = value
      updateWidgets(value)
    }
  })
  widget.value = widgetValue
  return { widget, minWidth, minHeight }
}

export const dynamicWidgets = { COMFY_DYNAMICCOMBO_V3: dynamicComboWidget }
const dynamicInputs: Record<
  string,
  (node: LGraphNode, inputSpec: InputSpecV2) => void
> = {
  COMFY_AUTOGROW_V3: applyAutogrow,
  COMFY_MATCHTYPE_V3: applyMatchType
}

export function applyDynamicInputs(
  node: LGraphNode,
  inputSpec: InputSpecV2
): boolean {
  if (!(inputSpec.type in dynamicInputs)) return false
  //TODO: move parsing/validation of inputSpec here?
  dynamicInputs[inputSpec.type](node, inputSpec)
  return true
}
function spliceInputs(
  node: LGraphNode,
  startIndex: number,
  deleteCount = -1,
  ...toAdd: INodeInputSlot[]
): INodeInputSlot[] {
  if (deleteCount < 0) return node.inputs.splice(startIndex)
  const ret = node.inputs.splice(startIndex, deleteCount, ...toAdd)
  node.inputs.slice(startIndex).forEach((input, index) => {
    const link = input.link && node.graph?.links?.get(input.link)
    if (link) link.target_slot = startIndex + index
  })
  return ret
}

function changeOutputType(
  node: LGraphNode,
  output: INodeOutputSlot,
  combinedType: ISlotType
) {
  if (output.type === combinedType) return
  output.type = combinedType

  //check and potentially remove links
  if (!node.graph) return
  for (const link_id of output.links ?? []) {
    const link = node.graph.links[link_id]
    if (!link) continue
    const { input, inputNode, subgraphOutput } = link.resolve(node.graph)
    const inputType = (input ?? subgraphOutput)?.type
    if (!inputType) continue
    const keep = LiteGraph.isValidConnection(combinedType, inputType)
    if (!keep && subgraphOutput) subgraphOutput.disconnect()
    else if (!keep && inputNode) inputNode.disconnectInput(link.target_slot)
    if (input && inputNode?.onConnectionsChange)
      inputNode.onConnectionsChange(
        LiteGraph.INPUT,
        link.target_slot,
        keep,
        link,
        input
      )
  }
}

function withComfyMatchType(node: LGraphNode): asserts node is MatchTypeNode {
  if (node.comfyMatchType) return
  node.comfyMatchType = {}

  const outputGroups = node.constructor.nodeData?.output_matchtypes
  node.onConnectionsChange = useChainCallback(
    node.onConnectionsChange,
    function (
      this: MatchTypeNode,
      contype: ISlotType,
      slot: number,
      iscon: boolean,
      linf: LLink | null | undefined
    ) {
      const input = this.inputs[slot]
      if (contype !== LiteGraph.INPUT || !this.graph || !input) return
      const [matchKey, matchGroup] = Object.entries(this.comfyMatchType).find(
        ([, group]) => input.name in group
      ) ?? ['', undefined]
      if (!matchGroup) return
      if (iscon && linf) {
        const { output, subgraphInput } = linf.resolve(this.graph)
        const connectingType = (output ?? subgraphInput)?.type
        if (connectingType) linf.type = connectingType
      }
      //NOTE: inputs contains input
      const groupInputs: INodeInputSlot[] = node.inputs.filter(
        (inp) => inp.name in matchGroup
      )
      const connectedTypes = groupInputs.map((inp) => {
        if (!inp.link) return '*'
        const link = this.graph!.links[inp.link]
        if (!link) return '*'
        const { output, subgraphInput } = link.resolve(this.graph!)
        return (output ?? subgraphInput)?.type ?? '*'
      })
      //An input slot can accept a connection that is
      // - Compatible with original type
      // - Compatible with all other input types
      //An output slot can output
      // - Only what every input can output
      groupInputs.forEach((input, idx) => {
        const otherConnected = [
          ...connectedTypes.slice(0, idx),
          ...connectedTypes.slice(idx + 1)
        ]
        const combinedType = commonType(
          ...otherConnected,
          matchGroup[input.name]
        )
        if (!combinedType) throw new Error('invalid connection')
        input.type = combinedType
      })
      const outputType = commonType(...connectedTypes)
      if (!outputType) throw new Error('invalid connection')
      this.outputs.forEach((output, idx) => {
        if (!(outputGroups?.[idx] == matchKey)) return
        changeOutputType(this, output, outputType)
      })
      app.canvas?.setDirty(true, true)
    }
  )
}

function applyMatchType(node: LGraphNode, inputSpec: InputSpecV2) {
  const { addNodeInput } = useLitegraphService()
  const name = inputSpec.name
  const { allowed_types, template_id } = (
    inputSpec as InputSpecV2 & {
      template: { allowed_types: string; template_id: string }
    }
  ).template
  const typedSpec = { ...inputSpec, type: allowed_types }
  addNodeInput(node, typedSpec)
  withComfyMatchType(node)
  node.comfyMatchType[template_id] ??= {}
  node.comfyMatchType[template_id][name] = allowed_types

  //TODO: instead apply on output add?
  //ensure outputs get updated
  const index = node.inputs.length - 1
  const input = node.inputs.at(-1)!
  requestAnimationFrame(() =>
    node.onConnectionsChange(LiteGraph.INPUT, index, false, undefined, input)
  )
}

function applyAutogrow(node: LGraphNode, untypedInputSpec: InputSpecV2) {
  const { addNodeInput } = useLitegraphService()

  const parseResult = zAutogrowOptions.safeParse(untypedInputSpec)
  if (!parseResult.success) throw new Error('invalid Autogrow spec')
  const inputSpec = parseResult.data

  const { input, min, names, prefix, max } = inputSpec.template
  const inputTypes: [Record<string, InputSpec> | undefined, boolean][] = [
    [input.required, false],
    [input.optional, true]
  ]
  const inputsV2 = inputTypes.flatMap(([inputType, isOptional]) =>
    Object.entries(inputType ?? {}).map(([name, v]) =>
      transformInputSpecV1ToV2(v, { name, isOptional })
    )
  )

  function nameToInputIndex(name: string) {
    const index = node.inputs.findIndex((input) => input.name === name)
    if (index === -1) throw new Error('Failed to find input')
    return index
  }
  function nameToInput(name: string) {
    return node.inputs[nameToInputIndex(name)]
  }

  //In the distance, someone shouting YAGNI
  const trackedInputs: string[][] = []
  function addInputGroup(insertionIndex: number) {
    const ordinal = trackedInputs.length
    const inputGroup = inputsV2.map((input) => ({
      ...input,
      name: names
        ? names[ordinal]
        : ((inputsV2.length == 1 ? prefix : input.name) ?? '') + ordinal,
      isOptional: ordinal >= (min ?? 0) || input.isOptional
    }))
    const newInputs = inputGroup
      .filter(
        (namedSpec) => !node.inputs.some((inp) => inp.name === namedSpec.name)
      )
      .map((namedSpec) => {
        addNodeInput(node, namedSpec)
        const input = spliceInputs(node, node.inputs.length - 1, 1)[0]
        if (inputsV2.length !== 1) ensureWidgetForInput(node, input)
        return input
      })
    spliceInputs(node, insertionIndex, 0, ...newInputs)
    trackedInputs.push(inputGroup.map((inp) => inp.name))
    app.canvas?.setDirty(true, true)
  }
  for (let i = 0; i < (min || 1); i++) addInputGroup(node.inputs.length)
  function removeInputGroup(inputName: string) {
    const groupIndex = trackedInputs.findIndex((ig) =>
      ig.some((inpName) => inpName === inputName)
    )
    if (groupIndex == -1) throw new Error('Failed to find group')
    const group = trackedInputs[groupIndex]
    for (const nameToRemove of group) {
      const inputIndex = nameToInputIndex(nameToRemove)
      const input = spliceInputs(node, inputIndex, 1)[0]
      if (!input.widget?.name) continue
      const widget = node.widgets?.find((w) => w.name === input.widget!.name)
      if (!widget) return
      widget.value = undefined
      node.removeWidget(widget)
    }
    trackedInputs.splice(groupIndex, 1)
    node.size[1] = node.computeSize([...node.size])[1]
    app.canvas?.setDirty(true, true)
  }

  function inputConnected(index: number) {
    const input = node.inputs[index]
    const groupIndex = trackedInputs.findIndex((ig) =>
      ig.some((inputName) => inputName === input.name)
    )
    if (groupIndex == -1) throw new Error('Failed to find group')
    if (
      groupIndex + 1 === trackedInputs.length &&
      trackedInputs.length < (max ?? names?.length ?? 100)
    ) {
      const lastInput = trackedInputs[groupIndex].at(-1)
      if (!lastInput) return
      const insertionIndex = nameToInputIndex(lastInput) + 1
      if (insertionIndex === 0) throw new Error('Failed to find Input')
      addInputGroup(insertionIndex)
    }
  }
  function inputDisconnected(index: number) {
    const input = node.inputs[index]
    if (trackedInputs.length === 1) return
    const groupIndex = trackedInputs.findIndex((ig) =>
      ig.some((inputName) => inputName === input.name)
    )
    if (groupIndex == -1) throw new Error('Failed to find group')
    if (
      trackedInputs[groupIndex].some(
        (inputName) => nameToInput(inputName).link != null
      )
    )
      return
    if (groupIndex + 1 < (min ?? 0)) return
    //For each group from here to last group, bubble swap links
    for (let column = 0; column < trackedInputs[0].length; column++) {
      let prevInput = nameToInputIndex(trackedInputs[groupIndex][column])
      for (let i = groupIndex + 1; i < trackedInputs.length; i++) {
        const curInput = nameToInputIndex(trackedInputs[i][column])
        const linkId = node.inputs[curInput].link
        node.inputs[prevInput].link = linkId
        const link = linkId && node.graph?.links?.[linkId]
        if (link) link.target_slot = prevInput
        prevInput = curInput
      }
      node.inputs[prevInput].link = null
    }
    if (
      trackedInputs.at(-2) &&
      !trackedInputs.at(-2)?.some((name) => !!nameToInput(name).link)
    )
      removeInputGroup(trackedInputs.at(-1)![0])
  }

  let pendingConnection: number | undefined
  let swappingConnection = false
  const originalOnConnectInput = node.onConnectInput
  node.onConnectInput = function (slot: number, ...args) {
    pendingConnection = slot
    requestAnimationFrame(() => (pendingConnection = undefined))
    return originalOnConnectInput?.apply(this, [slot, ...args]) ?? true
  }
  node.onConnectionsChange = useChainCallback(
    node.onConnectionsChange,
    (
      type: ISlotType,
      index: number,
      iscon: boolean,
      linf: LLink | null | undefined
    ) => {
      if (type !== NodeSlotType.INPUT) return
      const inputName = node.inputs[index].name
      if (!trackedInputs.flat().some((name) => name === inputName)) return
      if (iscon) {
        if (swappingConnection || !linf) return
        inputConnected(index)
      } else {
        if (pendingConnection === index) {
          swappingConnection = true
          requestAnimationFrame(() => (swappingConnection = false))
          return
        }
        requestAnimationFrame(() => inputDisconnected(index))
      }
    }
  )
}
